
import os
import json
import argparse
import sys 
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
sys.path.append("../../")
sys.path.append("..")
defaul_fontfamily="Times New Roman"
plot_setting_dict = {
    "dpi":300,
    "fig_size":(8,4),
    "line_width":1,
    "markersize":2,
    "label_fontdict":{
        'family' : defaul_fontfamily, 
        'size'   : 12
    },
    "ticks_fontdict":{
        'family' : defaul_fontfamily, 
        'size'   : 7
    },
    "title_fontdict":{
        'family' : defaul_fontfamily, 
        'size'   : 9
    },
    "legend_fontdict":{
        'family' : defaul_fontfamily, 
        'size'   : 8
    }
}

def getjsondata(casepath):
    result_list = []
    basename = os.path.basename(casepath)
    with open(casepath, "r") as fo:
        text = fo.read()
        jsondata = json.loads(text)
        print(f"Load {basename} data.")
    for i in jsondata:
        timedic = {}
        timedic["np"] = int(i.split("_")[-1])
        for j in jsondata[i]:
            if j == "Allreduce":
                timedic[j] = float(jsondata[i][j]["tavg"])
            elif j == "MYAllreduce":
                timedic["Three-Level-Allreduce"] = float(jsondata[i][j]["tavg"])
        result_list.append(timedic)
    return result_list

def myplot(savedir,dict):
    df = pd.DataFrame(dict)
    df = df.sort_index(axis=1)
    df = df.sort_values(by='np', ascending=True)
    df.set_index(["np"], inplace=True)
    
    # df = df/100000
    
    # fig, ax1 = plt.subplots(sharex=True,figsize=plot_setting_dict["fig_size"])
    # fig.add_subplot(111, frameon=False)

    # df.plot(ax=ax1, kind='bar', width=0.4)

    # # ax1.set_xticks(df.index)
    # # ax1.set_xticklabels(df.index, fontdict=plot_setting_dict["ticks_fontdict"], minor=False)

    # df["speedup"] = df["Allreduce"]/df["Three-Level-Allreduce"]
    # ax2 = ax1.twinx()
    # df["speedup"].plot(kind="line", ax=ax2, color='r', marker='o')

    # ax1.set_xlabel('Number of processor')
    # ax1.set_ylabel("Time(s)")
    # ax2.set_ylabel("Speedup")
    # ax1.grid(linestyle = '--')
    # # plt.title("")
    # plt.show()

def mycsv(savedir,dict):
    df = pd.DataFrame(dict)
    df = pd.DataFrame(dict)
    df = df.sort_index(axis=1)
    df = df.sort_values(by='np', ascending=True)
    df.set_index(["np"], inplace=True)
    
    writer = pd.ExcelWriter(os.path.join(savedir, "allreduce.xlsx"))
    df.to_excel(writer)
    writer.save()

def main():
    runcases = "/home/mrz/data/new_allreduce/ok_data/230425T2_useinamg/tianhe3_testpdl/JXPAMG_allreduce_test-RHD-1T.matrix-database.json"
    savedir = "/home/mrz/data/new_allreduce/ok_data/dataprocess0510"
    # print(runcases)
    datadict = getjsondata(runcases)
    # print(datadict)
    # print(list(datadict["jxpamg_solver04"]["np"]))
    # myplot(savedir, datadict)
    mycsv(savedir, datadict)
    # mycsv(args.savedir,datadict)

if __name__ == "__main__":
    main()